package demo

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
)

[4字节总长度] [4字节协议ID] [N字节协议体]
-----------------------------------------
|  length  |  protocolID |     data     |
|  (uint32)|   (uint32)  | (protobuf)   |
-----------------------------------------

// pkg/codec/length_header_codec.go
package codec

import (
"bytes"
"encoding/binary"
"errors"
"github.com/golang/protobuf/proto"
"io"
)

var ErrIncompletePacket = errors.New("incomplete packet")

type LengthHeaderCodec struct {
	TypeResolver func(uint32) proto.Message
}

// 编码流程：计算总长度 -> 写入头部 -> 写入protobuf数据
func (c *LengthHeaderCodec) Encode(protocolID uint32, msg proto.Message) ([]byte, error) {
	data, err := proto.Marshal(msg)
	if err != nil {
		return nil, err
	}

	totalLen := 4 + 4 + len(data) // 协议头8字节+数据长度
	buf := bytes.NewBuffer(make([]byte, 0, totalLen))

	// 写入长度头(4字节)
	if err := binary.Write(buf, binary.BigEndian, uint32(totalLen)); err != nil {
		return nil, err
	}

	// 写入协议ID(4字节)
	if err := binary.Write(buf, binary.BigEndian, protocolID); err != nil {
		return nil, err
	}

	// 写入协议体
	buf.Write(data)
	return buf.Bytes(), nil
}

// 解码流程：读取长度头 -> 读取完整包 -> 解析协议ID -> 解析protobuf
func (c *LengthHeaderCodec) Decode(r io.Reader) (uint32, proto.Message, error) {
	// 读取长度头
	lenBuf := make([]byte, 4)
	if _, err := io.ReadFull(r, lenBuf); err != nil {
		return 0, nil, err
	}

	totalLen := binary.BigEndian.Uint32(lenBuf)
	if totalLen < 8 {
		return 0, nil, errors.New("invalid packet length")
	}

	// 读取完整数据包
	packet := make([]byte, totalLen)
	copy(packet[:4], lenBuf) // 已经读取了4字节长度头

	if _, err := io.ReadFull(r, packet[4:]); err != nil {
		return 0, nil, err
	}

	// 解析协议ID
	protocolID := binary.BigEndian.Uint32(packet[4:8])

	// 获取消息类型
	msgType := c.TypeResolver(protocolID)
	if msgType == nil {
		return protocolID, nil, fmt.Errorf("unknown protocol id: %d", protocolID)
	}

	// 解析协议体
	if err := proto.Unmarshal(packet[8:], msgType); err != nil {
		return protocolID, nil, err
	}

	return protocolID, msgType, nil
}